// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/common/base/exceptions.h>

#include <melon/container/f14_map.h>

namespace kumo::pollux {
    /// A priority queue with values of type 'T'. Each value has assigned priority
    /// on insertion which determines the value's position in the queue. It supports
    /// to update the priority of the existing values, which adjusts the value's
    /// position in the queue accordingly. Ties are broken by insertion and update
    /// order. If 'kMaxQueue' is true, it is a max priority queue, otherwise a min
    /// priority queue.
    ///
    /// NOTE: this class is not thread-safe.
    template<
        typename T,
        bool kMaxQueue,
        typename Allocator = std::allocator<T>,
        typename Hash = std::hash<T>,
        typename EqualTo = std::equal_to<T> >
    class IndexedPriorityQueue {
    public:
        explicit IndexedPriorityQueue(const Allocator &allocator = {})
            : values_(allocator),
              priorities_(RebindAlloc<int64_t>(allocator)),
              generations_(RebindAlloc<int64_t>(allocator)),
              heap_(RebindAlloc<int32_t>(allocator)),
              valueIndices_(RebindAlloc<std::pair<const T, int32_t> >(allocator)),
              heapIndices_(RebindAlloc<int32_t>(allocator)) {
        }

        ~IndexedPriorityQueue() {
            validate();
        }

        int32_t size() const {
            return values_.size();
        }

        bool empty() const {
            return values_.empty();
        }

        /// Inserts 'value' into the queue with specified 'priority'. If 'value'
        /// already exists, then update its priority and the corresponding position in
        /// the queue.
        bool addOrUpdate(const T &value, int64_t priority) {
            auto it = valueIndices_.find(value);
            if (it == valueIndices_.end()) {
                addNewValue(value, priority);
                return true;
            }
            auto i = it->second;
            if (priorities_[i] == priority) {
                return false;
            }
            updatePriority(i, priority);
            return true;
        }

        const T &top() const {
            POLLUX_DCHECK(!heap_.empty());
            return values_[heap_[0]];
        }

        int64_t topPriority() const {
            POLLUX_DCHECK(!heap_.empty());
            return priorities_[heap_[0]];
        }

        T pop() {
            POLLUX_DCHECK(!heap_.empty());
            auto i = heap_[0];
            heap_[0] = heap_.back();
            heapIndices_[heap_.back()] = 0;
            heap_.pop_back();
            if (!heap_.empty()) {
                percolateDown(0);
            }
            auto result = values_[i];
            valueIndices_.erase(values_[i]);
            if (i != size() - 1) {
                valueIndices_[values_.back()] = i;
                values_[i] = values_.back();
                priorities_[i] = priorities_.back();
                generations_[i] = generations_.back();
                heap_[heapIndices_.back()] = i;
                heapIndices_[i] = heapIndices_.back();
            }
            values_.pop_back();
            priorities_.pop_back();
            generations_.pop_back();
            heapIndices_.pop_back();
            return result;
        }

        const T *values() const {
            return values_.data();
        }

        const int64_t *priorities() const {
            return priorities_.data();
        }

        std::optional<int> getValueIndex(const T &value) const {
            auto it = valueIndices_.find(value);
            if (it != valueIndices_.end()) {
                return it->second;
            }
            return std::nullopt;
        }

        void updatePriority(int index, int64_t priority) {
            bool up = priority < priorities_[index];
            if constexpr (kMaxQueue) {
                up = !up;
            }
            priorities_[index] = priority;
            generations_[index] = ++generation_;
            if (up) {
                percolateUp(heapIndices_[index]);
            } else {
                percolateDown(heapIndices_[index]);
            }
        }

        void addNewValue(const T &value, int64_t priority) {
            POLLUX_DCHECK_LT(size(), std::numeric_limits<int32_t>::max());
            auto i = size();
            values_.push_back(value);
            priorities_.push_back(priority);
            generations_.push_back(++generation_);
            POLLUX_CHECK(valueIndices_.emplace(value, i).second);
            heapIndices_.push_back(i);
            heap_.push_back(i);
            percolateUp(i);
        }

        void replaceTop(const T &value, int64_t priority) {
            POLLUX_DCHECK(!heap_.empty());
            int i = heap_[0];
            valueIndices_.erase(values_[i]);
            values_[i] = value;
            priorities_[i] = priority;
            generations_[i] = ++generation_;
            POLLUX_CHECK(valueIndices_.emplace(value, i).second);
            percolateDown(0);
        }

        /// Return negative number if value at i should be ordered before j; positive
        /// number if j should be ordered before i. Otherwise return 0.
        int64_t compare(int i, int j) const {
            int64_t result = priorities_[i] - priorities_[j];
            if constexpr (kMaxQueue) {
                result = -result;
            }
            return result != 0 ? result : generations_[i] - generations_[j];
        }

    private:
        template<typename U>
        using RebindAlloc =
        typename std::allocator_traits<Allocator>::template rebind_alloc<U>;

        void validate() const {
            POLLUX_DCHECK_EQ(values_.size(), priorities_.size());
            POLLUX_DCHECK_EQ(values_.size(), generations_.size());
            POLLUX_DCHECK_EQ(values_.size(), heap_.size());
            POLLUX_DCHECK_EQ(values_.size(), valueIndices_.size());
            POLLUX_DCHECK_EQ(values_.size(), heapIndices_.size());
        }

        void percolateUp(int pos) {
            while (pos > 0) {
                int parent = (pos - 1) / 2;
                if (compare(heap_[pos], heap_[parent]) >= 0) {
                    break;
                }
                std::swap(heap_[pos], heap_[parent]);
                heapIndices_[heap_[pos]] = pos;
                pos = parent;
            }
            heapIndices_[heap_[pos]] = pos;
        }

        void percolateDown(int pos) {
            for (;;) {
                int left = 2 * pos + 1;
                if (left >= heap_.size()) {
                    break;
                }
                int right = left + 1;
                int next = right < heap_.size() && compare(heap_[right], heap_[left]) < 0
                               ? right
                               : left;
                if (compare(heap_[pos], heap_[next]) <= 0) {
                    break;
                }
                std::swap(heap_[pos], heap_[next]);
                heapIndices_[heap_[pos]] = pos;
                pos = next;
            }
            heapIndices_[heap_[pos]] = pos;
        }

        int64_t generation_{0};
        std::vector<T, Allocator> values_;
        std::vector<int64_t, RebindAlloc<int64_t> > priorities_;
        std::vector<int64_t, RebindAlloc<int64_t> > generations_;
        std::vector<int32_t, RebindAlloc<int32_t> > heap_;
        melon::F14FastMap<
            T,
            int32_t,
            Hash,
            EqualTo,
            RebindAlloc<std::pair<const T, int32_t> > >
        valueIndices_;
        std::vector<int32_t, RebindAlloc<int32_t> > heapIndices_;
    };
} // namespace kumo::pollux
